"""使用FasterRCNN ONNX进行目标检测，ONNX下载地址：
https://github.com/onnx/models/blob/master/vision/object_detection_segmentation/faster-rcnn/model/FasterRCNN-10.onnx
"""

import cv2
import numpy as np
import onnxruntime


class VehicleFasterRCNN(object):
    def __init__(self, onnx_file_path: str, threshold: float = 0.5):
        self.session = onnxruntime.InferenceSession(
            onnx_file_path,
            providers=["CUDAExecutionProvider", "CPUExecutionProvider"])
        # TODO: 默认输入大小 1920 x 1080，此处resize成32的倍数，960x544
        # 如果输入大小发生改变，此处需要调整
        self.detect_size = (960, 544)
        # 参考：
        # https://github.com/onnx/models/blob/master/vision/object_detection_segmentation/faster-rcnn/dependencies/coco_classes.txt
        # 与车辆相关的indices是（3，6，8），base 0
        # 3: car, 6: bus, 8: truck
        # self.coco_vehicle_indices = (3, 6, 8)
        self.coco_vehicle_indices = (3,)
        self.threshold = threshold

    def _preprocess(self, image: np.array) -> np.array:
        image = cv2.resize(image, dsize=self.detect_size)
        image = np.array(image, dtype=np.float32)
        image = image.transpose(2, 0, 1)
        mean_vec = np.array([102.9801, 115.9465, 122.7717])
        mean_vec = mean_vec[:, None, None]
        image -= mean_vec
        return image

    def _post_process(self, boxes, labels, scores):
        keep = np.zeros(shape=(len(labels), ), dtype="bool")
        for index in self.coco_vehicle_indices:
            keep |= (labels == index)
        keep = keep & (scores > self.threshold)
        boxes = boxes[keep]
        labels = labels[keep]
        scores = scores[keep]
        boxes[:, [0, 2]] /= self.detect_size[0]
        boxes[:, [1, 3]] /= self.detect_size[1]
        w = boxes[:, 2] - boxes[:, 0]
        h = boxes[:, 3] - boxes[:, 1]
        keep = (w < 0.5) & (h < 0.5)
        boxes = boxes[keep]
        return boxes

    def detect(self, image: np.array) -> np.ndarray:
        """检车图像中车辆的位置

        Parameters
        ----------
        image : np.array
            输入的BGR图像

        Returns
        -------
        np.ndarray
            所有车辆的位置, [x1, y1, x2, y2]格式
        """
        image = self._preprocess(image)
        session_outputs = self.session.run(None, {"image": image})
        boxes = self._post_process(*session_outputs)
        return boxes
